import pika
from common.json_helper import dict_to_json
from conf.config import Config


class AMQPServer(object):
    def __init__(self, **kwargs):
        virtual_host = Config.amqp.virtual_host
        if kwargs.get('virtual_host'):
            virtual_host = kwargs['virtual_host']

        userx = pika.PlainCredentials(Config.amqp.user_name, Config.amqp.password)
        self.connection = pika.BlockingConnection(
            pika.ConnectionParameters(
                host=Config.amqp.host,
                virtual_host=virtual_host,
                credentials=userx)
        )

    def __enter__(self):
        return self.connection

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.connection.close()


def publish_message(
        message_id,
        message,
        exchange=Config.amqp.exchange,
        routing_key=Config.amqp.routing_key,
        virtual_host=Config.amqp.virtual_host,
        **kwargs):

    if not exchange:
        exchange = Config.amqp.exchange
    if not routing_key:
        routing_key = Config.amqp.routing_key
    if not virtual_host:
        virtual_host = Config.amqp.virtual_host

    with AMQPServer(**kwargs) as amqp:
        channel = amqp.channel()
        if type(message) == dict:
            message = dict_to_json(message)

        channel.basic_publish(
            exchange=exchange,
            routing_key=routing_key,
            body=message,
            properties=pika.BasicProperties(
                message_id=message_id,
                # delivery_mode=2 # 可以指定此条消息持久化，防止RabbitMQ服务挂掉之后消息丢失但是此属性设置并不能百分百保证消息真的被持久化，因为RabbitMQ挂掉的时候它可能还保存在缓存中，没来得及同步到磁盘中
            )
        )
        print(f" [x] Sent message to host is {virtual_host},exchange is {exchange}, routing key is {routing_key}")
        print(" [x] message id %r message is %r" % (message_id, message))
    return virtual_host, exchange, routing_key


def consume_message(on_message_callback, prefetch_count=1, **kwargs):
    with AMQPServer(**kwargs) as amqp:
        print(f" [x] host is {Config.amqp.virtual_host}, queue is {Config.amqp.queue}")
        channel = amqp.channel()
        channel.basic_qos(prefetch_count=prefetch_count)
        channel.basic_consume(queue=Config.amqp.queue, on_message_callback=on_message_callback)
        print('Waiting for messages.')
        channel.start_consuming()
